# -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
import numpy as np

from libs.configs import cfgs
from help_utils.densely_coded_label import angle_label_decode, get_code_len


def angle_cls_focal_loss(labels, pred, anchor_state, alpha=None, gamma=2.0, decimal_weight=False):

    indices = tf.reshape(tf.where(tf.equal(anchor_state, 1)), [-1, ])
    labels = tf.gather(labels, indices)
    pred = tf.gather(pred, indices)
    anchor_state = tf.gather(anchor_state, indices)

    # compute the focal loss
    per_entry_cross_ent = - labels * tf.log(tf.sigmoid(pred) + cfgs.EPSILON) \
                          - (1 - labels) * tf.log(1 - tf.sigmoid(pred) + cfgs.EPSILON)

    prediction_probabilities = tf.sigmoid(pred)
    p_t = ((labels * prediction_probabilities) +
           ((1 - labels) * (1 - prediction_probabilities)))
    modulating_factor = 1.0
    if gamma:
        modulating_factor = tf.pow(1.0 - p_t, gamma)
    alpha_weight_factor = 1.0
    if alpha is not None:
        alpha_weight_factor = (labels * alpha +
                               (1 - labels) * (1 - alpha))

    if decimal_weight:
        angle_decode_labels = tf.py_func(func=angle_label_decode,
                                         inp=[labels, cfgs.ANGLE_RANGE, cfgs.OMEGA, cfgs.ANGLE_MODE],
                                         Tout=[tf.float32])
        angle_decode_labels = tf.reshape(angle_decode_labels, [-1, ]) * -1

        angle_decode_pred = tf.py_func(func=angle_label_decode,
                                       inp=[tf.sigmoid(pred), cfgs.ANGLE_RANGE, cfgs.OMEGA, cfgs.ANGLE_MODE],
                                       Tout=[tf.float32])

        angle_decode_pred = tf.reshape(angle_decode_pred, [-1, ]) * -1

        diff_weight = tf.reshape(tf.log(abs(angle_decode_labels - angle_decode_pred) + 1), [-1, 1])
    else:
        diff_weight = tf.ones_like(tf.reshape(anchor_state, [-1, 1]))

    focal_cross_entropy_loss = (diff_weight * modulating_factor * alpha_weight_factor *
                                per_entry_cross_ent)

    # compute the normalizer: the number of positive anchors
    # normalizer = tf.stop_gradient(tf.where(tf.greater(anchor_state, -2)))
    normalizer = tf.stop_gradient(tf.where(tf.equal(anchor_state, 1)))
    normalizer = tf.cast(tf.shape(normalizer)[0], tf.float32)
    normalizer = tf.maximum(1.0, normalizer)

    # normalizer = tf.stop_gradient(tf.cast(tf.equal(anchor_state, 1), tf.float32))
    # normalizer = tf.maximum(tf.reduce_sum(normalizer), 1)

    return tf.reduce_sum(focal_cross_entropy_loss) / normalizer


def angle_cls_period_focal_loss(labels, pred, anchor_state, target_boxes, alpha=None, gamma=2.0,
                                decimal_weight=False, aspect_ratio_threshold=1.5):

    indices = tf.reshape(tf.where(tf.equal(anchor_state, 1)), [-1, ])
    labels = tf.gather(labels, indices)
    pred = tf.gather(pred, indices)
    target_boxes = tf.gather(target_boxes, indices)
    anchor_state = tf.gather(anchor_state, indices)

    # compute the focal loss
    per_entry_cross_ent = - labels * tf.log(tf.sigmoid(pred) + cfgs.EPSILON) \
                          - (1 - labels) * tf.log(1 - tf.sigmoid(pred) + cfgs.EPSILON)

    prediction_probabilities = tf.sigmoid(pred)
    p_t = ((labels * prediction_probabilities) +
           ((1 - labels) * (1 - prediction_probabilities)))
    modulating_factor = 1.0
    if gamma:
        modulating_factor = tf.pow(1.0 - p_t, gamma)
    alpha_weight_factor = 1.0
    if alpha is not None:
        alpha_weight_factor = (labels * alpha +
                               (1 - labels) * (1 - alpha))

    if decimal_weight:
        angle_decode_labels = tf.py_func(func=angle_label_decode,
                                         inp=[labels, cfgs.ANGLE_RANGE, cfgs.OMEGA, cfgs.ANGLE_MODE],
                                         Tout=[tf.float32])
        angle_decode_labels = tf.reshape(angle_decode_labels, [-1, ]) * -1

        angle_decode_pred = tf.py_func(func=angle_label_decode,
                                       inp=[tf.sigmoid(pred), cfgs.ANGLE_RANGE, cfgs.OMEGA, cfgs.ANGLE_MODE],
                                       Tout=[tf.float32])

        angle_decode_pred = tf.reshape(angle_decode_pred, [-1, ]) * -1

        target_boxes = tf.reshape(target_boxes[:, :-1], [-1, 5])
        x, y, h, w, theta = tf.unstack(target_boxes, axis=-1)
        aspect_ratio = h / w
        period_weight_90 = tf.cast(tf.less_equal(aspect_ratio, aspect_ratio_threshold), tf.int32) * 2 * 180 / cfgs.ANGLE_RANGE
        period_weight_180 = tf.cast(tf.greater(aspect_ratio, aspect_ratio_threshold), tf.int32) * 1 * 180 / cfgs.ANGLE_RANGE

        period_weight = tf.cast(period_weight_90 + period_weight_180, tf.float32)
        diff_weight = tf.reshape(tf.abs(tf.sin(period_weight * (angle_decode_labels - angle_decode_pred))), [-1, 1])

    else:
        diff_weight = tf.ones_like(tf.reshape(anchor_state, [-1, 1]))

    focal_cross_entropy_loss = (diff_weight * modulating_factor * alpha_weight_factor *
                                per_entry_cross_ent)

    # compute the normalizer: the number of positive anchors
    # normalizer = tf.stop_gradient(tf.where(tf.greater(anchor_state, -2)))
    normalizer = tf.stop_gradient(tf.where(tf.equal(anchor_state, 1)))
    normalizer = tf.cast(tf.shape(normalizer)[0], tf.float32)
    normalizer = tf.maximum(1.0, normalizer)

    # normalizer = tf.stop_gradient(tf.cast(tf.equal(anchor_state, 1), tf.float32))
    # normalizer = tf.maximum(tf.reduce_sum(normalizer), 1)

    return tf.reduce_sum(focal_cross_entropy_loss) / normalizer

